
import cv2
import numpy as np
class PositionProjection:

    def __init__(self, lmsk_x, lmsk_y):
        """
        计算仿射变换的类
        将网球场一个画面的点转换到另外一个画面
        lmsk_x :  画面x的关键点，只取其中观点的几个点，网球场有18个关键点，取其中4个
        lmsk_y :  画面y的关键点，只取其中观点的几个点
        """
        key_lmsk_index = {1,2,16,17}
        standard_lmks = {1:[421,560], 2:[1244,560], 16:[421,2937], 17:[1244,2937]}
        lmsk_x_key_pos = []
        lmsk_y_key_pos = []
        standard_key_pos = []
        for i in key_lmsk_index:
            lmsk_x_key_pos.append(lmsk_x[str(i)][:2])
            lmsk_y_key_pos.append(lmsk_y[str(i)][:2])
            standard_key_pos.append(standard_lmks[i])

        # 计算仿射变换矩阵
        self.m_x2std = cv2.getPerspectiveTransform(np.array(lmsk_x_key_pos, dtype=np.float32), np.array(standard_key_pos, dtype=np.float32))
        self.m_std2x = cv2.getPerspectiveTransform(np.array(standard_key_pos, dtype=np.float32), np.array(lmsk_x_key_pos, dtype=np.float32))

        self.m_y2std = cv2.getPerspectiveTransform(np.array(lmsk_y_key_pos, dtype=np.float32), np.array(standard_key_pos, dtype=np.float32))
        self.m_std2y = cv2.getPerspectiveTransform(np.array(standard_key_pos, dtype=np.float32), np.array(lmsk_y_key_pos, dtype=np.float32))

    def transform_x_to_y(self, pos_x):
        '''
        将画面x的点转换成画面y的点
        '''
        pos_x = np.array([pos_x], dtype=np.float32)
        pos_std = cv2.perspectiveTransform(pos_x[None, :], self.m_x2std)[0][0]
        pos_std = np.array([[1665 - pos_std[0], 3506 - pos_std[1]]], dtype=np.float32)
        pos_y = cv2.perspectiveTransform(pos_std[None, :], self.m_std2y)[0][0]
        pos_y = [int(pos_y[0]), int(pos_y[1])]
        return pos_y

    def transform_y_to_x(self, pos_y):
        '''
        将画面y的点转换成画面x的点
        '''
        pos_y = np.array([pos_y], dtype=np.float32)
        pos_std = cv2.perspectiveTransform(pos_y[None, :], self.m_y2std)[0][0]
        pos_std = np.array([[1665 - pos_std[0], 3506 - pos_std[1]]], dtype=np.float32)
        pos_x = cv2.perspectiveTransform(pos_std[None, :], self.m_std2x)[0][0]
        pos_x = [int(pos_x[0]), int(pos_x[1])]
        return pos_x
    
def plot_lmks(img, lmks):
    """
    绘制关键点
    """
    for key,value in lmks.items():
        cv2.circle(img, (value[0], value[1]), 2, (0, 255, 0), -1)
        cv2.putText(img, str(key), (value[0], value[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    return img
    

def mouse_callback(event, x, y, flags, param):
    """
    鼠标回调函数
    """
    if event == cv2.EVENT_LBUTTONDOWN:
        print("x:{}, y:{}".format(x, y))
        img_x = param[0]
        img_y = param[1]
        lmks_x = param[2]
        lmks_y = param[3]
        # 计算仿射变换矩阵
        pos_proj = PositionProjection(lmks_x, lmks_y)
        pos_y = pos_proj.transform_x_to_y([x, y])
        print("pos_x:{}, pos_y:{}".format([x, y], pos_y))
        cv2.circle(img_y, (pos_y[0], pos_y[1]), 2, (0, 255, 0), -1)
        cv2.circle(img_x, (x, y), 2, (0, 255, 0), -1)
        cv2.imshow("img_y", img_y)

if __name__ == '__main__':
    # 读取图片
    path_x = "/Users/xclin/Desktop/output_image.jpg"
    path_y = "/Users/xclin/Desktop/output_image2.jpg"
    img_x = cv2.imread(path_x)
    img_y = cv2.imread(path_y)
    # 读取关键点,关键点的坐标可以通过鼠标来输出
    lmks_x = {"1":[558,234], "2":[753,234], "16":[276,515], "17":[1031,515]}
    lmks_y = {"1":[541,188], "2":[744,184], "16":[238,542], "17":[1058,529]}
    # 绘制关键点
    img_x = plot_lmks(img_x, lmks_x)
    img_y = plot_lmks(img_y, lmks_y)
    # 显示图片
    cv2.namedWindow("img_x", cv2.WINDOW_NORMAL)
    cv2.imshow("img_x", img_x)
    cv2.namedWindow("img_y", cv2.WINDOW_NORMAL)
    cv2.setMouseCallback("img_x", mouse_callback, param=[img_x, img_y, lmks_x, lmks_y])
    cv2.imshow("img_y", img_y)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 转换点
    # pos_x = [1000, 1000]
    # pos_y = pos_proj.transform_x_to_y(pos_x)
    # print("pos_x:{}, pos_y:{}".format(pos_x, pos_y))
    # pos_x = pos_proj.transform_y_to_x(pos_y)
    # print("pos_x:{}, pos_y:{}".format(pos_x, pos_y))
